import os

import numpy as np
from torchvision import datasets

from .custom import CustomClsDataset
from ..common.types import *


class AdapterCustomClsDataset(CustomClsDataset):
    """Adapter class for built-in datasets in torchvision.
    
    - MNIST
    - CIFAR10
    """

    MAPPING = {
        "mnist": datasets.MNIST,
        "cifar10": datasets.CIFAR10
    }

    def __init__(self, name: str, **kwargs):
        assert name.lower() in self.MAPPING, "%s is not supported!" % name
        self.name = name
        super().__init__(root=os.path.expanduser("~/dataset"), **kwargs)

    def load(self):
        self.dataset = self.MAPPING[self.name](root=os.path.expanduser("~/dataset"),
                                               train=self.partition, download=True)

    def __len__(self) -> int:
        return len(self.dataset)

    def __getitem__(self, index: int) -> ClsDataItem:
        image, label = self.dataset[index]
        image = np.array(image)

        return image, label
